#include <cassert>
#include <iostream>
#include <random>
#include <vector>
#include <algorithm>
#include <functional>
#include "bcstatistics.hpp"

using std::cout;
using std::endl;
using std::vector;

typedef unsigned char uchar;
typedef std::independent_bits_engine<std::mt19937,8,uchar> byte_generator;

namespace {

void operator^=(vector<uchar> & io, vector<uchar> const& rhs)
{
	assert(io.size() == rhs.size());
	for (unsigned i=0,e=io.size(); i<e; ++i) {
		io[i] ^= rhs[i];
	}
}

void bit_counter(vector<long> & c, vector<uchar> const& b)
{
	const unsigned b2cs = b.size()*8;
	if (c.size() < b2cs)
		c.resize(b2cs);
	auto cit = c.begin();
	for (unsigned i=0,e=b.size(); i<e; ++i) {
		unsigned byte = b[i];
		for (int bit=0; bit<8; ++bit) {
			*cit += ((byte >> bit) & 1);
			++cit;
		}
	}
}

void plain_text_difference(blockcipher& bc, byte_generator& reng, int pt_bit_index, long trials)
{
	vector<uchar> key (bc.key_bytes());
	vector<uchar> blk1 (bc.block_bytes());
	vector<uchar> blk2;
	vector<long> bitcounters;
	assert(pt_bit_index < bc.block_bytes()*8);
	for (long t=0; t<trials; ++t) {
		std::generate(key.begin(),key.end(),std::ref(reng));
		std::generate(blk1.begin(),blk1.end(),std::ref(reng));
		blk2 = blk1;
		blk2[pt_bit_index/8] ^= (1 << (pt_bit_index%8));
		bc.keysetup(key.data());
		bc.encrypt(blk1.data());
		bc.encrypt(blk2.data());
		blk1 ^= blk2;
		bit_counter(bitcounters,blk1);
	}
	cout << "\rAfter " << trials << " iterations for plaintext bit ";
	cout.width(3);
	cout << pt_bit_index << ": ";
	int mii = 0;
	int maa = 0;
	for (unsigned i=0; i<bitcounters.size(); ++i) {
		int prom = (bitcounters[i]*1000+trials/2) / trials;
		if (i==0 || prom<mii) mii=prom;
		if (i==0 || prom>maa) maa=prom;
	}
	cout << mii/10 << "." << mii%10 << "% ... ";
	cout << maa/10 << "." << maa%10 << "%\n";
}

} // anonymous namespace

void run_bc_test(blockcipher& bc)
{
	cout << "[-BLOCKCIPHER STATISTICAL TESTS-]\n"
		"block cipher: "; bc.identify(cout); cout << "\n"
		"block length: " << bc.block_bytes() << " bytes\n"
		"key   length: " << bc.key_bytes() << " bytes\n";
	byte_generator bg;
	for (int bit=0; bit<bc.block_bytes()*8; ++bit)
		plain_text_difference(bc,bg,bit,333000);
}
